# -*- coding: utf-8 -*-
from mazedata import *
from scipy.sparse.csgraph import floyd_warshall

# 下标转为三角形
freetri = tri.simplices[freetri]

# 三角形的质心
center = np.zeros((freetri.shape[0], 2))
for i in range(3):
    center += points[freetri[:, i]]
center /= 3

trimax = 0
for i in range(freetri.shape[0]):
    for j in range(3):
        d = np.linalg.norm(center[i] - points[freetri[i][j]])
        if trimax < d:
            trimax = d

# 质心的KD树版
kdcenter = KDTree(center)

def lst_sorted(lst):
    ix = min(lst)
    ix = lst.index(ix)
    return tuple(lst[ix:] + lst[:ix])

edge2tri = dict()

for it, lt in enumerate(freetri):
    for i in range(-1, 2):
        edge = lst_sorted([lt[i], lt[i+1]])
        if edge not in edge2tri:
            edge2tri[edge] = []
        edge2tri[edge].append(it)

tri2edge = dict()
obstacle_edge = set()

for edge, trilst in edge2tri.items():
    if len(trilst) == 2:
        tri2edge[lst_sorted(trilst)] = edge
    else:
        obstacle_edge.add(edge)

distance = np.zeros((freetri.shape[0], freetri.shape[0]))


for trilst, edge in tri2edge.items():
    mid = points[edge[0]] + points[edge[1]]
    mid /= 2
    d = np.linalg.norm(center[trilst[0]] - mid)
    d += np.linalg.norm(center[trilst[1]] - mid)
    distance[trilst[0], trilst[1]] = d
    distance[trilst[1], trilst[0]] = d

shorted, predecessors = floyd_warshall(distance, return_predecessors=True)

def IsInTri(idxlst, pos):
    p = points[idxlst]
    mat = np.zeros((3, 3))
    mat[:, 2] = 1.0
    for i in range(-1, 2):
        mat[0, :2] = p[i]
        mat[1, :2] = p[i + 1]
        mat[2, :2] = pos
        if np.linalg.det(mat) < 0:
            return False
    return True

def GetTriByPos(pos):
    trilst = kdcenter.query_ball_point(pos, trimax+0.1)
    for t in trilst:
        if IsInTri(freetri[t], pos):
            return t
    return kdcenter.query(pos)[1]


def GetPath(here, there):
    here_t = GetTriByPos(here)
    there_t = GetTriByPos(there)
    if not IsInTri(freetri[there_t], there):
        there = center[there_t]
    trilst = [here_t]
    while True:
        next_tri = predecessors[there_t, trilst[-1]]
        if next_tri < 0:
            break
        if next_tri != there_t:
            trilst.append(next_tri)
        else:
            trilst.append(there_t)
            break
    pos_lst = [here]
    for i in range(len(trilst) - 1):
        t1 = trilst[i]
        t2 = trilst[(i+1)]
        edge = tri2edge[lst_sorted([t1, t2])]
        mid = (points[edge[0]] + points[edge[1]]) / 2
        pos_lst.append(mid)
    pos_lst.append(there)
    return np.array(pos_lst)


def GetPathDistance(here, there):
    path = GetPath(here, there)
    dis = 0.0
    for i in range(path.shape[0]-1):
        dis += np.linalg.norm(path[i] - path[i+1])
    return dis


def RandomPos():
    pos = np.random.random((2,))
    pos[0] *= WIDTH
    pos[1] *= HEIGHT
    _, ix = kdcenter.query(pos)
    return kdcenter.data[ix]